// Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved

#pragma once

#include <cstdint>
#include <memory>
#include <c10/core/Device.h>
#include <stddef.h>
#include <stdbool.h>

bool kpex_use_hbm();

int kpex_posix_memalign(void **memptr, size_t alignment, size_t size);

void kpex_free(void *ptr);

namespace kpex {
inline c10::Device device() {
    return c10::Device(c10::kCPU, (int)kpex_use_hbm());
}

template <typename T>
struct KpexMallocDeleter {
    void operator()(T *ptr) const
    {
        kpex_free(ptr);
    }
};

template <typename T>
inline std::unique_ptr<T[], KpexMallocDeleter<T> > alloc(int64_t size)
{
    void *ptr;
    kpex_posix_memalign(&ptr, 64, size * sizeof(T));
    return std::unique_ptr<T[], KpexMallocDeleter<T> >((T *)ptr);
}
}   // namespace kpex